package info.wangyuan.agent.aop;

import info.wangyuan.agent.exception.BusinessException;
import info.wangyuan.agent.utils.JwtUtils;
import jakarta.servlet.http.HttpServletRequest;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.time.Duration;

/**
 * 根据 OpenId 对话限流切面
 *
 * @author Albert
 * @since 2025-09-11 20:18:24
 */
@Aspect
@Component
public class OpenIdRateLimitAspect {

    private final StringRedisTemplate redisTemplate;
    private final JwtUtils jwtUtils;

    public OpenIdRateLimitAspect(StringRedisTemplate redisTemplate, JwtUtils jwtUtils) {
        this.redisTemplate = redisTemplate;
        this.jwtUtils = jwtUtils;
    }

    private static final String OpenId_LIMIT_KEY_PREFIX = "openId:chatLimit:";

    @Around("@annotation(openIdRateLimit)")
    public Object around(ProceedingJoinPoint joinPoint, OpenIdRateLimit openIdRateLimit) throws Throwable {
        // 获取请求的IP
        ServletRequestAttributes attrs = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (attrs == null) return joinPoint.proceed();
        HttpServletRequest request = attrs.getRequest();

        // 这时已经过了拦截器那一层，所以 token 请求头必然存在
        String aiToken = request.getHeader("AiToken");
        String openid = jwtUtils.getOpenidFromToken(aiToken);

        // 获取设置的限制参数
        int limit = openIdRateLimit.limit();
        long period = openIdRateLimit.period();

        String key = OpenId_LIMIT_KEY_PREFIX + openid;

        // Redis计数
        Long count = redisTemplate.opsForValue().increment(key);
        if (count != null && count == 1L) {
            redisTemplate.expire(key, Duration.ofSeconds(period));
        }

        if (count != null && count > limit) {
            throw new BusinessException("请求过于频繁，请稍后重试");
        }

        return joinPoint.proceed();
    }
}
